{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Rotation as scipy_R\n",
    "from typing import Union, List\n",
    "from shapely.geometry import LineString, Polygon\n",
    "import pickle\n",
    "# with open('../tmp/result/result_aux_470.pkl','rb') as f:\n",
    "#     data = pickle.load(f)\n",
    "\n",
    "def init_axes(dpi: float=100) -> plt.Axes:\n",
    "    fig = plt.figure(dpi=dpi)\n",
    "    return fig.add_axes([0, 0, 1, 1])\n",
    "\n",
    "def matplot_2d_cloud(cloud_x, cloud_y, scale=0.1, rgb='b',\n",
    "                     alpha=1, label=None, axes: plt.Axes=None) -> None:\n",
    "    \"\"\" plot 2D cloud with matplotlib\n",
    "    :param cloud_x: list of cloud x\n",
    "    :param cloud_y: list of cloud y\n",
    "    :param scale: size of cloud points, float or list of float\n",
    "    :param rgb: color used for plot\n",
    "    :param alpha: opacity of line, 0 for transparent, 1 for opaque\n",
    "    :param label: label of cloud\n",
    "    :param axes: axes handler of matplotlib\n",
    "    \"\"\"\n",
    "    if axes is None:\n",
    "        axes = init_axes(dpi=1000)\n",
    "    if label is None:\n",
    "        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none')\n",
    "    else:\n",
    "        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none', label=label)\n",
    "    axes.set_xlabel('x axis')\n",
    "    axes.set_ylabel('y axis')\n",
    "    axes.axis('equal')\n",
    "    return axes\n",
    "def matplot_polygon(points_x, points_y, scale=0.1, rgb='g',\n",
    "                    alpha=1, axes:plt.Axes=None) -> None:\n",
    "    if axes is None:\n",
    "        axes = init_axes()\n",
    "    assert len(points_x) == len(points_y) and len(points_x) > 1\n",
    "    x_arr = np.append(np.array(points_x), points_x[0])\n",
    "    y_arr = np.append(np.array(points_y), points_y[0])\n",
    "    axes.plot(x_arr, y_arr, c=rgb, lw=scale, alpha=alpha)\n",
    "    axes.set_xlabel('x axis')\n",
    "    axes.set_ylabel('y axis')\n",
    "    axes.axis('equal')\n",
    "    return axes\n",
    "def matplot_lines(points_x, points_y, scale=0.1, rgb='g',\n",
    "                    alpha=1, axes:plt.Axes=None) -> None:\n",
    "    if axes is None:\n",
    "        axes = init_axes()\n",
    "    assert len(points_x) == len(points_y) and len(points_x) > 1\n",
    "    x_arr = np.array(points_x)\n",
    "    y_arr = np.array(points_y)\n",
    "    axes.plot(x_arr, y_arr, c=rgb, lw=scale, alpha=alpha)\n",
    "    axes.set_xlabel('x axis')\n",
    "    axes.set_ylabel('y axis')\n",
    "    axes.axis('equal')\n",
    "    axes\n",
    "    return axes\n",
    "\n",
    "\n",
    "DEFAULT_ALL_CLOSE_PRECISION = 1.e-4\n",
    "\n",
    "def euler_to_rotation_matrix(euler_xyz: Union[np.ndarray, list], degree=False) -> np.ndarray:\n",
    "    return scipy_R.from_euler('XYZ', euler_xyz, degrees=degree).as_matrix()\n",
    "\n",
    "def bounding_box_to_points(position: list, size: list, rotation: list) -> np.ndarray:\n",
    "    \"\"\" calculate 3D bounding box's points\n",
    "    :param position: position of box\n",
    "    :param size: size of box (x, y, z)\n",
    "    :param rotation: heading of box in euler angles\n",
    "    \"\"\"\n",
    "    assert len(position) == len(size) == len(rotation) == 3\n",
    "    pos_arr = np.array(position)\n",
    "    size_arr = np.array(size)\n",
    "    pts_prefix = np.array([[-1, -1, -1], [-1, 1, -1], [1, 1, -1], [1, -1, -1],\n",
    "                           [-1, -1, 1], [-1, 1, 1], [1, 1, 1], [1, -1, 1]]) / 2\n",
    "    return euler_to_rotation_matrix(rotation).dot((pts_prefix * size_arr).T).T + pos_arr\n",
    "def matplot_bounding_box(position: list, size: list, rotation: list, scale=0.1,\n",
    "                         rgb='k', alpha=1, axes: plt.Axes=None) -> None:\n",
    "    \"\"\" plot 3D bounding box with matplotlib\n",
    "    :param position: position of box\n",
    "    :param size: size of box (x, y, z)\n",
    "    :param rotation: heading of box in euler angles\n",
    "    :param scale: line width of box\n",
    "    :param rgb: color used for plot\n",
    "    :param alpha: opacity of line, 0 for transparent, 1 for opaque\n",
    "    :param axes: axes handler of matplotlib\n",
    "    \"\"\"\n",
    "    if axes is None:\n",
    "        axes = init_axes()\n",
    "    box_pts = bounding_box_to_points(position, size, rotation)[:5, :2]\n",
    "    axes.plot(box_pts[:, 0], box_pts[:, 1], c=rgb, lw=scale, alpha=alpha)\n",
    "    axes.set_xlabel('x axis')\n",
    "    axes.set_ylabel('y axis')\n",
    "    axes.axis('equal')\n",
    "    return axes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_norm2(a, b):\n",
    "    return np.linalg.norm(a - b, axis=-1, ord=2)\n",
    "def get_next_points(pts: np.array, sorted_pts: list, eps=0.2):\n",
    "    if sorted_pts is None:\n",
    "        sorted_pts = [pts[0][0], pts[0][1]]\n",
    "        pts = pts[1:]\n",
    "    pt = sorted_pts[-1:]\n",
    "    \n",
    "    d = []\n",
    "    d.append(get_norm2(pts[:,0,:], pt))\n",
    "    d.append(get_norm2(pts[:,1,:], pt))\n",
    "    # d.append(get_norm2(pts[:,0,:], pt[0,:]))\n",
    "    # d.append(get_norm2(pts[:,1,:], pt[0,:]))\n",
    "    # d.append(get_norm2(pts[:,0,:], pt[1,:]))\n",
    "    # d.append(get_norm2(pts[:,1,:], pt[1,:]))\n",
    "    d = np.stack(d, -1)\n",
    "    # idx = np.argmin(d, -1)\n",
    "    mind = np.min(d)\n",
    "    flag = 0\n",
    "    if mind < eps:\n",
    "        indx = np.where(d == mind)\n",
    "\n",
    "        cur_pts_idx = indx[1] // 2\n",
    "        fut_pts_idx = indx[1] % 2\n",
    "        # print(cur_pts_idx,fut_pts_idx)\n",
    "        # import pdb;pdb.set_trace()\n",
    "        # print(indx[0][0])\n",
    "        next_pts = pts[indx[0][0]].reshape(2, 2)\n",
    "        pts = np.delete(pts, indx[0], axis=0)\n",
    "\n",
    "        # print(next_pts,fut_pts_idx)\n",
    "        join_pts = next_pts[fut_pts_idx[0]]\n",
    "        next_pt = next_pts[1-fut_pts_idx[0]]\n",
    "        # print(join_pts, pt, next_pt, next_pts)\n",
    "\n",
    "        # fusion join pt\n",
    "        sorted_pts[-1] = (join_pts + pt).reshape(-1) / 2\n",
    "\n",
    "        sorted_pts.append(next_pt.reshape(-1))\n",
    "    else:\n",
    "        flag = 1\n",
    "    return sorted_pts, pts, flag\n",
    "    \n",
    "def get_sorted_lines(pts1,eps=0.2):\n",
    "    lines_pts_sets = []\n",
    "    sorted_pts = None\n",
    "    while 1:\n",
    "        # print(len(lines_pts_sets), pts1.shape[0])\n",
    "        if len(pts1) > 1:\n",
    "            sorted_pts, pts1, flag = get_next_points(pts1, sorted_pts,eps=eps)\n",
    "            if flag:\n",
    "                lines_pts_sets.append(sorted_pts)\n",
    "                sorted_pts = None\n",
    "        else:\n",
    "            break\n",
    "    return lines_pts_sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [

    "def result_plot(data,frame=0, pred_map=True,pred_box=False,gt_map=False, lidar=True,gt_box=False, map_score=0.5,box_score=0.2,decode_map=False):\n",
    "    rgb = ['r','y','b','g']\n",
    "    scale = 1\n",
    "    axes = init_axes(dpi=250)\n",
    "    if pred_map:\n",
    "        for i in range(len(data[frame]['maps'])):\n",
    "            cls = data[frame]['maps'][i]['maps_class'].cpu().squeeze(0).numpy()\n",
    "            pts = data[frame]['maps'][i]['maps_points'].cpu().squeeze(0).numpy()\n",
    "            scr = data[frame]['maps'][i]['maps_scores'].cpu().squeeze(0).numpy()\n",
    "            mask = scr > map_score\n",
    "            pts1 = pts[mask]\n",
    "            axes = matplot_2d_cloud(pts1[:,0,0], pts1[:,0,1],scale=0.5,rgb=rgb[i],label=0,axes=axes)\n",
    "            axes = matplot_2d_cloud(pts1[:,1,0], pts1[:,1,1],scale=0.5,rgb=rgb[i],label=1,axes=axes)\n",
    "    if gt_map:\n",
    "        gt_maps = data[frame]['gt_maps_3d'][0][0].cpu().numpy()\n",
    "        for i in range(int(gt_maps[:,0].max())):\n",
    "            gt_map = gt_maps[gt_maps[:,0]==i,1:]\n",
    "            if len(gt_map)<=0:\n",
    "                continue\n",
    "            class_= int(gt_map[0,-1])\n",
    "            axes = matplot_lines(gt_map[:,0], gt_map[:,1],rgb=rgb[class_],scale=0.3,axes=axes)\n",
    "    # lidar\n",
    "    if lidar and data[frame].get('points', False):\n",
    "        points = data[frame]['points'][0].cpu().numpy()\n",
    "        axes = matplot_2d_cloud(points[:,0], points[:,1],scale=0.1,rgb='g',label=4,axes=axes)\n",
    "\n",
    "    # draw bbox\n",
    "    if pred_box:\n",
    "        box = data[frame][\"pts_bbox\"]['boxes_3d']\n",
    "        mask = data[frame]['pts_bbox']['scores_3d'] > box_score\n",
    "        box = box[mask]\n",
    "        box_center = box.center.cpu().numpy().tolist()\n",
    "        box_dims = box.dims.cpu().numpy().tolist()\n",
    "        box_yaw = box.yaw.cpu().numpy().tolist()\n",
    "        for b in range(len(box_yaw)):\n",
    "            axes = matplot_bounding_box(box_center[b],box_dims[b],[0,0,box_yaw[b]],axes=axes, scale=0.4,rgb='r')\n",
    "    if gt_box and data[frame].get('gt_bboxes_3d', False):\n",
    "        box = data[frame]['gt_bboxes_3d'][0][0]\n",
    "        # box = data[frame][\"pts_bbox\"]['boxes_3d']\n",
    "        # mask = data[frame]['pts_bbox']['scores_3d'] > box_score\n",
    "        # box = box[mask]\n",
    "        box_center = box.center.cpu().numpy().tolist()\n",
    "        box_dims = box.dims.cpu().numpy().tolist()\n",
    "        box_yaw = box.yaw.cpu().numpy().tolist()\n",
    "        for b in range(len(box_yaw)):\n",
    "            axes = matplot_bounding_box(box_center[b],box_dims[b],[0,0,box_yaw[b]],axes=axes,scale=0.5,rgb='k')\n",
    "    plt.ylim(-60, 60)\n",
    "    plt.xlim(-60, 60)\n",
    "# gt_points = []\n",
    "# for l in data_gt['infos'][0]['vec_lines']['lane_divider']:\n",
    "#     line = np.array(l['points'])\n",
    "#     axes = matplot_lines(line[:,0], line[:,1],rgb='y',scale=0.3,axes=axes)\n",
    "# for l in data_gt['infos'][0]['vec_lines']['road_edge']:\n",
    "#     line = np.array(l['points'])\n",
    "#     axes = matplot_lines(line[:,0], line[:,1],rgb='r',scale=0.3,axes=axes)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  7%|▋         | 20/300 [00:25<06:11,  1.33s/it]/tmp/ipykernel_3434782/703049609.py:12: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`). Consider using `matplotlib.pyplot.close()`.\n",
      "  fig = plt.figure(dpi=dpi)\n",
      "100%|█████████▉| 299/300 [03:53<00:01,  1.29s/it]"
     ]
    }
   ],
   "source": [
    "# frame = np.random.randint(low=0,high=len(data))\n",
    "# frame =0\n",
    "# result_plot(data, frame,lidar=False, map_score=0.45, box_score=0.45,  pred_box=True, pred_map=False, gt_map=False,gt_box=False,decode_map=True)\n",
    "from tqdm import tqdm, trange\n",
    "for frame in tqdm(range(300)):\n",
    "    try:\n",
    "        result_plot(data, frame,lidar=False, map_score=0.45, box_score=0.45,  pred_box=True, pred_map=False, gt_map=False,gt_box=False,decode_map=True)\n",
    "    except:\n",
    "        print(frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# frame =  np.random.randint(low=0,high=len(data))\n",
    "# frame=200\n",
    "result_plot(data, frame, map_score=0.50, box_score=0.1,  pred_box=False, pred_map=True, gt_map=True, gt_box=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "video = cv2.VideoWriter('../tmp/test.avi',cv2.VideoWriter_fourcc(*'MJPG'),10,(1280,720))\n",
    "\n",
    "for i in range(300):\n",
    "    #读取图片\n",
    "    img = cv2.imread(f'../tmp/video/{i:06d}.png')     \n",
    "   \t# resize方法是cv2库提供的更改像素大小的方法\n",
    "    # 将图片转换为1280*720像素大小\n",
    "    img = cv2.resize(img,(1280,720))\n",
    "    # 写入视频\n",
    "    video.write(img)\n",
    "    print(i)\n",
    "# 释放资源\n",
    "video.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data[480]['timestamp'][0].item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# frame = np.random.randint(low=0,high=len(data))\n",
    "# frame = 1455\n",
    "def save_pig(frame):\n",
    "    rgb = ['r','y','b','g']\n",
    "    lw = [0.5,0.5,0.5,0.3]\n",
    "    axes = init_axes(dpi=1000)\n",
    "    for i in range(4):\n",
    "        cls = data[frame]['maps'][i]['maps_class'].cpu().squeeze(0).numpy()\n",
    "        pts = data[frame]['maps'][i]['maps_points'].cpu().squeeze(0).numpy()\n",
    "        scr = data[frame]['maps'][i]['maps_scores'].cpu().squeeze(0).numpy()\n",
    "        mask = scr > 0.5\n",
    "        pts1 = pts[mask]\n",
    "        idx = np.argsort(pts1[:,0, 0])# + pts1[:,0, 1])\n",
    "        pts1 = pts1[idx]\n",
    "        lines = get_sorted_lines(pts1, eps=0.2)\n",
    "        for k, l in enumerate(lines):\n",
    "            l = np.stack(l,0)\n",
    "            # if len(l) < 3:\n",
    "            #     continue\n",
    "            # print(k)\n",
    "            axes = matplot_lines(l[:,0], l[:,1], rgb=rgb[i], scale=lw[i], axes=axes)\n",
    "    axes.set_xbound(-50,50)\n",
    "    axes.set_ybound(-50,50)\n",
    "    # plt.savefig(f'../tmp/video/{frame:06d}.png')\n",
    "save_pig(frame)\n",
    "# save_pig(66)\n",
    "# for i in range(300):\n",
    "#     try:\n",
    "#         save_pig(i)\n",
    "#     except:\n",
    "#         pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame = 100\n",
    "rgb = ['r','y','b']\n",
    "axes = init_axes(dpi=1000)\n",
    "for i in range(3):\n",
    "    cls = data[frame]['maps'][i]['maps_class'].cpu().squeeze(0).numpy()\n",
    "    pts = data[frame]['maps'][i]['maps_points'].cpu().squeeze(0).numpy()\n",
    "    scr = data[frame]['maps'][i]['maps_scores'].cpu().squeeze(0).numpy()\n",
    "    mask = scr > 0.4\n",
    "    pts1 = pts[mask]\n",
    "    idx = np.argsort(pts1[:,0, 0])# + pts1[:,0, 1])\n",
    "    pts1 = pts1[idx]\n",
    "    lines = get_sorted_lines(pts1,eps=0.4)\n",
    "    for l in lines:\n",
    "        l = np.stack(l,0)\n",
    "        if len(l) < 3:\n",
    "            continue\n",
    "        axes = matplot_lines(l[:,0], l[:,1], rgb=np.random.rand(3).tolist(), scale=0.5, axes=axes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Rotation as scipy_R\n",
    "from typing import Union, List                                \n",
    "from shapely.geometry import LineString, Polygon\n",
    "import pickle\n",
    "with open('../tmp/result/result_argov2_4d_val_gap1_2000.pkl','rb') as f:\n",
    "    data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data[0].keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open('../tmp/result/result_aux_map_gt.pkl','wb') as f:\n",
    "#     pickle.dump(data,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i in data:#[195]['pts_bbox']\n",
    "#     i.pop(\"maps\")\n",
    "#     # i.pop(\"gt_maps_3d\")\n",
    "    # pred_box = i.pop('pts_bbox')\n",
    "    # i['pred_box'] = pred_box['boxes_3d'].tensor.cpu().numpy()\n",
    "    # i['pred_labels'] = pred_box['labels_3d'].cpu().numpy()\n",
    "    # i['pred_scores'] = pred_box['scores_3d'].cpu().numpy()\n",
    "    # i['gt_bboxes_3d'] = i['gt_bboxes_3d'][0][0].tensor.cpu().numpy()\n",
    "    # i['gt_labels_3d'] = i['gt_labels_3d'][0][0].cpu().numpy()\n",
    "    # i['timestamp'] = np.int(i['timestamp'][0].cpu().numpy())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i in data:\n",
    "#     i.pop('points')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# frame = data[195]\n",
    "\n",
    "# frame['maps'][1][0]['heatmap']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# frame = data[195]\n",
    "# reg = frame['maps'][1][0]['reg'].cpu().numpy()\n",
    "# scores = frame['maps'][1][0]['heatmap'].sigmoid().cpu().numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# for check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Rotation as scipy_R\n",
    "from typing import Union, List\n",
    "from shapely.geometry import LineString, Polygon\n",
    "import pickle\n",
    "with open('../tmp/result/result_165.pkl','rb') as f:\n",
    "    data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_multi_grid_range(axes, grid_size=0.8, bv_range=(-51.2, -51.2, 51.2, 51.2),rgb='k'):\n",
    "    Lx, Ly, Lz, Lx1, Ly1, Lz1 = [], [], [], [], [], []\n",
    "\n",
    "    Lx = np.arange(bv_range[0], bv_range[2] + 0.1, grid_size) #bugfix\n",
    "    Ly1 = np.arange(bv_range[1], bv_range[3] + 0.1, grid_size)\n",
    "    Lx = np.stack([Lx, Lx], axis=0).transpose().reshape(-1)\n",
    "    Ly1 = np.stack([Ly1, Ly1], axis=0).transpose().reshape(-1)\n",
    "\n",
    "    for i in range(int(Lx.shape[0] / 2)):\n",
    "        if (i + 1) % 2 == 0:\n",
    "            Ly.append(bv_range[1])\n",
    "            Ly.append(bv_range[1 + 2])\n",
    "        else:\n",
    "            Ly.append(bv_range[1 + 2])\n",
    "            Ly.append(bv_range[1])\n",
    "    Ly = np.array(Ly)\n",
    "    Lz = np.zeros(Ly.shape[0])\n",
    "\n",
    "    for i in range(int(Ly1.shape[0] / 2)):\n",
    "        if (i + 1) % 2 == 0:\n",
    "            Lx1.append(bv_range[0])\n",
    "            Lx1.append(bv_range[0 + 2])\n",
    "        else:\n",
    "            Lx1.append(bv_range[0 + 2])\n",
    "            Lx1.append(bv_range[0])\n",
    "    # print(Lx)\n",
    "    Lx1 = np.array(Lx1)\n",
    "    Lz1 = np.zeros(Lx1.shape[0])\n",
    "\n",
    "    axes = matplot_lines(Lx1, Ly1, rgb=rgb, scale=0.05, axes=axes)\n",
    "    axes = matplot_lines(Lx, Ly, rgb=rgb, scale=0.05, axes=axes)\n",
    "    # Lx = np.concatenate([Lx,L])\n",
    "    return axes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame = np.random.randint(low=0,high=len(data))\n",
    "# frame=211\n",
    "rgb = ['r','y','b','g']\n",
    "def get_rgb(r):\n",
    "    if len(r) == 0:\n",
    "        return []\n",
    "    l = r.min()\n",
    "    u = r.max()\n",
    "    r = r - l\n",
    "    r = (r / (u - l))\n",
    "    rb = np.zeros((r.shape[0],3))\n",
    "    rb[:,0] = r\n",
    "    rb[:,2] = 1-r\n",
    "    return rb\n",
    "axes = init_axes(dpi=1000)\n",
    "for i in range(4):\n",
    "    cls = data[frame]['maps'][i]['maps_class'].cpu().squeeze(0).numpy()\n",
    "    pts = data[frame]['maps'][i]['maps_points'].cpu().squeeze(0).numpy()\n",
    "    scr = data[frame]['maps'][i]['maps_scores'].cpu().squeeze(0).numpy()\n",
    "    mask = scr > 0.5\n",
    "    pts1 = pts[mask]\n",
    "    scr = scr[mask]\n",
    "    scr_rgb = get_rgb(scr)\n",
    "    # idx = np.argsort(pts1[:,0, 0])# + pts1[:,0, 1])\n",
    "    # pts1 = pts1[idx]\n",
    "    # lines = get_sorted_lines(pts1,eps=0.2)\n",
    "    for k, l in enumerate(pts1):\n",
    "        sc = scr[k]\n",
    "        if k > 10000:\n",
    "            break\n",
    "        # l = np.stack(l,0)\n",
    "        # if len(l) < 3:\n",
    "            # continue\n",
    "        # print(k)\n",
    "        axes = matplot_lines(l[:,0], l[:,1], rgb=rgb[i], scale=0.05, axes=axes)\n",
    "draw_multi_grid_range(axes)\n",
    "draw_multi_grid_range(axes,grid_size=0.8, bv_range=(-30.4, -15.2, 30.4, 15.2),rgb='g')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# for check gt gengerate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Rotation as scipy_R\n",
    "from typing import Union, List\n",
    "from shapely.geometry import LineString, Polygon\n",
    "import pickle\n",
    "import sys\n",
    "sys.path.pop(-2)\n",
    "from mmdet3d.ops.voxel_interpolation import voxel_interpolation_cpu as VI\n",
    "with open('../tmp/result/result_120.pkl','rb') as f:\n",
    "    data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def in_range(gt_dense_pts, pts_range):\n",
    "    pts_mask = np.logical_and(gt_dense_pts[:, 1] >= pts_range[0], gt_dense_pts[:, 1] <= pts_range[3])\n",
    "    pts_mask = np.logical_and(pts_mask, gt_dense_pts[:, 2] <= pts_range[4])\n",
    "    pts_mask = np.logical_and(pts_mask, gt_dense_pts[:, 2] >= pts_range[1])\n",
    "\n",
    "    return pts_mask\n",
    "\n",
    "def outlier_gt_filter(gt_dense_pts, pts_range):\n",
    "    lines_num = int(np.max(gt_dense_pts[:, 0]) + 1)\n",
    "\n",
    "    pts_mask = in_range(gt_dense_pts, pts_range)\n",
    "    new_gt_dense_pts = np.zeros(gt_dense_pts.shape, dtype=np.float32)\n",
    "\n",
    "    id_count = 0\n",
    "    filled_size = 0\n",
    "    for i in range(lines_num):\n",
    "        pos_mask = (gt_dense_pts[:, 0] == i)\n",
    "        if (pts_mask[pos_mask].sum() == 0):\n",
    "            continue\n",
    "\n",
    "        lines_gt = gt_dense_pts[pos_mask, :]\n",
    "        lines_gt[:, 0] = id_count\n",
    "        id_count += 1\n",
    "\n",
    "        circle = (np.abs(lines_gt[0, 1] - lines_gt[-1, 1])**2 + np.abs(lines_gt[0, 2] - lines_gt[-1, 2])**2) < 1\n",
    "        if(circle and\n",
    "            (pts_mask[pos_mask].sum() != pos_mask.sum()) and\n",
    "            in_range(lines_gt[0, :].reshape(1, -1), pts_range)[0]):     #环形 且 存在越界部分 且 起点在范围内\n",
    "            first_outside = np.where(~pts_mask[pos_mask])[0][0]\n",
    "            lines_gt = np.concatenate([lines_gt[first_outside: -1, :],\n",
    "                                        lines_gt[:first_outside, :],\n",
    "                                        lines_gt[first_outside, :].reshape(1, -1)])\n",
    "\n",
    "        new_gt_dense_pts[filled_size: filled_size + len(lines_gt), :] = lines_gt\n",
    "        filled_size += len(lines_gt)\n",
    "\n",
    "    new_gt_dense_pts = new_gt_dense_pts[:filled_size]\n",
    "    return new_gt_dense_pts\n",
    "def decode_one_voxels(dt_pts):\n",
    "    dx = (dt_pts <= 1) * 0 + \\\n",
    "            ((1 < dt_pts) & (dt_pts <= 2)) * (dt_pts - 1) + \\\n",
    "            ((2 < dt_pts) & (dt_pts <= 3)) * 1 + \\\n",
    "            (3 < dt_pts) * (4 - dt_pts)\n",
    "    dy = (dt_pts <= 1) * dt_pts + \\\n",
    "            ((1 < dt_pts) & (dt_pts <= 2)) * 1 + \\\n",
    "            ((2 < dt_pts) & (dt_pts <= 3)) * (3 - dt_pts) + \\\n",
    "            (3 < dt_pts) * 0\n",
    "    return dx, dy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame = np.random.randint(low=0,high=len(data))\n",
    "# import sys\n",
    "# sys.path.append(\"..\")\n",
    "from mmdet3d.ops.voxel_interpolation import voxel_interpolation_cpu as VI\n",
    "\n",
    "rgb = ['r','y','b']\n",
    "gt_maps = data[frame]['gt_maps_3d'][0][0]\n",
    "max_objs = 10000\n",
    "# pc_range = torch.tensor(self.train_cfg['point_cloud_range'])\n",
    "# voxel_size = torch.tensor(self.train_cfg['voxel_size'])\n",
    "\n",
    "# feature_map_size = grid_size[:2] // self.train_cfg['out_size_factor']\n",
    "vs = [0.8, 0.8, 8]\n",
    "pcr = [-51.2, -51.2, -5.0, 51.2, 51.2, 3.0]\n",
    "gt_pts = gt_maps.cpu().numpy()\n",
    "gt_pts = np.fromfile('../pts.bin',dtype=np.float32).reshape(-1, 5)\n",
    "gt_pts = gt_pts.astype(np.float32)\n",
    "gt_pts = outlier_gt_filter(gt_pts, pcr)\n",
    "dense_pts = np.zeros((max_objs, 8), dtype=np.float32)\n",
    "import time\n",
    "st = time.time()\n",
    "size = VI.gt_lines_voxel_interpolation(gt_pts, vs, pcr, \n",
    "                                    max_objs, dense_pts)\n",
    "dense_pts = dense_pts[:size, :]\n",
    "print(time.time() - st)\n",
    "dx1,dy1 = decode_one_voxels(dense_pts[:,5])\n",
    "dy2,dx2 = decode_one_voxels(dense_pts[:,6])\n",
    "xs = dense_pts[:,2] * vs[0] + pcr[0]\n",
    "ys = dense_pts[:,3] * vs[1] + pcr[1]\n",
    "pts = np.stack([dx1,dy1,dx2,dy2],-1).reshape(-1, 2, 2)\n",
    "pts *= vs[0]\n",
    "pts[...,0] += xs.reshape(-1,1)\n",
    "pts[...,1] += ys.reshape(-1,1)\n",
    "axes = init_axes(dpi=2000)\n",
    "draw_multi_grid_range(axes)\n",
    "draw_multi_grid_range(axes, grid_size=0.8, bv_range=(-30.4, -15.2, 30.4, 15.2), rgb='g')\n",
    "for k, l in enumerate(pts):\n",
    "    # if dense_pts[k,0] != 0:\n",
    "    #     break\n",
    "    sc = int(dense_pts[k][-1])\n",
    "    if k > 10000:\n",
    "        break\n",
    "    # print(l.shape)\n",
    "    # l = np.stack(l,0)\n",
    "    # if len(l) < 3:\n",
    "        # continue\n",
    "    # print(k)\n",
    "    axes = matplot_lines(l[:,0]+k*0.0001, l[:,1] + k*0.0001, rgb=rgb[sc], scale=0.05, axes=axes)\n",
    "    # x, y = l[0,0],l[0,1]\n",
    "    # dx, dy = l[1,0]-l[0,0],l[1,1] - l[0,1]\n",
    "    # axes.arrow(x, y, dx, dy,width=0.00005,head_width=0.00005)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_objs = 10000\n",
    "# pc_range = torch.tensor(self.train_cfg['point_cloud_range'])\n",
    "# voxel_size = torch.tensor(self.train_cfg['voxel_size'])\n",
    "\n",
    "# feature_map_size = grid_size[:2] // self.train_cfg['out_size_factor']\n",
    "vs = [0.8, 0.8, 8]\n",
    "pcr = [-51.2, -51.2, -5.0, 51.2, 51.2, 3.0]\n",
    "gt_pts = gt_maps.cpu().numpy()\n",
    "gt_pts = np.fromfile('../pts.bin',dtype=np.float32).reshape(-1, 5)\n",
    "gt_pts = gt_pts.astype(np.float32)\n",
    "gt_pts = outlier_gt_filter(gt_pts, pcr)\n",
    "dense_pts = np.zeros((max_objs, 8), dtype=np.float32)\n",
    "import time\n",
    "st = time.time()\n",
    "size = VI.gt_lines_voxel_interpolation(gt_pts, vs, pcr, \n",
    "                                    max_objs, dense_pts)\n",
    "dense_pts = dense_pts[:size, :]\n",
    "print(time.time() - st)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_pts\n",
    "for k, l in enumerate(pts):\n",
    "    # if dense_pts[k,0] != 0:\n",
    "    #     break\n",
    "    sc = int(dense_pts[k][-1])\n",
    "    if k > 10000:\n",
    "        break\n",
    "    # print(l.shape)\n",
    "    # l = np.stack(l,0)\n",
    "    # if len(l) < 3:\n",
    "        # continue\n",
    "    # print(k)\n",
    "    axes = matplot_lines(l[:,0]+k*0.0001, l[:,1] + k*0.0001, rgb=rgb[sc], scale=0.05, axes=axes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.pop(-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open('../data/argoverse2/sensor/bevdetv2-argo_infos_train_map.pkl','rb') as f:\n",
    "    data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data[\"infos\"][0].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, d in enumerate(data['infos']):\n",
    "    d['ego2global'] = d['cams']['ring_rear_left']['ego2global']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../data/argoverse2/sensor/bevdetv2-argo_infos_train_map.pkl','wb') as f:\n",
    "    pickle.dump(data, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open('../data/argoverse2/sensor/bevdetv2-argo_infos_val_map.pkl','rb') as f:\n",
    "    data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "for d in data['infos']:\n",
    "    for i in d['vec_lines']['road_markings']:\n",
    "        if i['category'][0] != 'crosswalk':\n",
    "            print(i['category'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'id': 0,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[124.96404618446569, 2.7587248575730863, -2.136524964072038],\n",
       "   [99.44991157773131, 2.547150769580307, -0.809057199353222],\n",
       "   [79.10735611947382, 2.566969632550454, -0.635522064750115],\n",
       "   [43.55218426017194, 2.5702096778965897, -0.6002460324601638],\n",
       "   [32.3942964542605, 2.2624645821272136, -0.5202178702643815],\n",
       "   [22.18162942808872, 1.8134470137503573, -0.4006488050070286],\n",
       "   [10.07878489217962, 1.3970602650956683, -0.34132264546890867],\n",
       "   [-29.373781493526145, 1.409962119081538, -0.11976587400849059],\n",
       "   [-33.48869557825668, 1.350768202597812, -0.1932299780049287],\n",
       "   [-41.51679692313926, 1.397618575420707, -0.12328020225863481],\n",
       "   [-58.2706781635028, 1.2654448137664407, -0.10937047439128285]]},\n",
       " {'id': 1,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[124.96020597549864, 3.138473274476481, -2.1498062537042735],\n",
       "   [99.45016694720212, 2.926913082063431, -0.8224930024192226],\n",
       "   [79.10746257862593, 2.9467322137220338, -0.6489522772541885],\n",
       "   [43.54967523523315, 2.94996745732786, -0.6135777148531237],\n",
       "   [32.38073728792284, 2.642003148000018, -0.5331262405455585],\n",
       "   [22.16680791170606, 2.192940045809577, -0.41350809235535346],\n",
       "   [10.075740970237735, 1.776814843406749, -0.3546340988500726],\n",
       "   [-29.379011139409226, 1.7896957851917739, -0.1329943896693322],\n",
       "   [-33.486597891423564, 1.7305223600774826, -0.20673477667289664],\n",
       "   [-41.51980391603195, 1.7773733989535085, -0.1365930530979025],\n",
       "   [-58.27368515639549, 1.645199637299244, -0.12268332523055077]]},\n",
       " {'id': 2,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[141.31971722335913, 14.291489884396542, -3.980435158725281],\n",
       "   [140.3677881285448, 25.769991982406356, -4.861128399547592]]},\n",
       " {'id': 3,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[140.74627582203678, 25.80020647385764, -4.876430372661109],\n",
       "   [141.69820491685113, 14.321704375847824, -3.9957371318387978]]},\n",
       " {'id': 4,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[140.86282388657241, -8.05932534661976, -2.962777660457761],\n",
       "   [140.61790760366645, -22.59672931871914, -3.8514855800635215]]},\n",
       " {'id': 5,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[140.23820903097734, -22.59119338040339, -3.837401997960987],\n",
       "   [140.4831253138833, -8.053789408304008, -2.9486940783552265]]},\n",
       " {'id': 6,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[138.50358959616986, 25.46728745397472, -4.810359654710077],\n",
       "   [136.88925757131238, 13.850963045379947, -3.8182712776972783]]},\n",
       " {'id': 7,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[136.51321469595464, 13.904268918517317, -3.806014043846116],\n",
       "   [138.12754672081212, 25.52059332711209, -4.798102420858916]]},\n",
       " {'id': 8,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[-80.895791591376, 15.698809277979572, 0.030944224869468507],\n",
       "   [-82.32078551701427, 25.868211886117976, -0.12479770580613393],\n",
       "   [-85.64987241590906, 50.0, -0.49550353074786424]]},\n",
       " {'id': 9,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[-80.51976816960872, 15.751254848642496, 0.014948963060951215],\n",
       "   [-81.94465016864595, 25.919856076802688, -0.140768844669385],\n",
       "   [-85.26661252302705, 50.0, -0.5106813259011144]]},\n",
       " {'id': 10,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'white',\n",
       "  'points': [[-41.77221210543853, 4.482287531182002, -0.2227308764510785],\n",
       "   [-58.284463372609714, 4.536685031271164, -0.19446425663940275]]},\n",
       " {'id': 11,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[-71.7164344678863, -50.0, 0.9907728725931803],\n",
       "   [-71.89158299639938, -48.719835750656614, 0.9684630796313916],\n",
       "   [-72.39964843418193, -45.30681642592598, 0.9269851106015184],\n",
       "   [-72.47155234811213, -44.65754408420013, 0.9167479932220949],\n",
       "   [-76.72444844634302, -12.589832281584885, 0.5137163293875165]]},\n",
       " {'id': 12,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[-72.09957769564717, -50.0, 1.005898623607022],\n",
       "   [-72.26748505291923, -48.77313653497151, 0.9844840126739913],\n",
       "   [-72.77552781557719, -45.360275460510806, 0.9430107857165694],\n",
       "   [-72.84794952002828, -44.70726277531432, 0.9326609039367516],\n",
       "   [-77.10084561825917, -12.639550972699075, 0.5296292401021742]]},\n",
       " {'id': 13,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[-92.09395942977018, 1.9167784773353316, 0.9092860123888946],\n",
       "   [-132.40966520955266, 2.003220807639008, 3.0932720932811684]]},\n",
       " {'id': 14,\n",
       "  'line_type': 'solid',\n",
       "  'color': 'yellow',\n",
       "  'points': [[-132.4095784497357, 2.382983419982344, 3.079842620506727],\n",
       "   [-92.09387266995326, 2.2965410896786747, 0.8958565396144538]]}]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['infos'][0]['vec_lines']['lane_dividers']"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.13 ('mm_zhl')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "52337d3598aa7e5bd30d88474e67ee45f9d9ba8e83642815fedffdbb63dec38f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
